#!/usr/bin/env python3
# -*- coding: utf-8 -*-

from flask import Flask, request, jsonify, render_template, send_from_directory
import json, re, random, string, datetime, ollama, os
from werkzeug.utils import secure_filename
# Import the database connector functions
from db_connector import create_test_case, update_test_case, test_connection, search_test_cases

app = Flask(__name__)
app.config['SECRET_KEY'] = ''.join(random.choices(string.ascii_letters + string.digits, k=32))
app.config['UPLOAD_FOLDER'] = 'uploads'
app.config['MAX_CONTENT_LENGTH'] = 16 * 1024 * 1024  # 16MB max upload

# 确保上传目录存在
os.makedirs(app.config['UPLOAD_FOLDER'], exist_ok=True)

# Check if database connection is available
DB_AVAILABLE = test_connection()
print(f"MySQL Database connection status: {'Available' if DB_AVAILABLE else 'Not available'}")

# 从main.py导入的函数
def clean_json_response(text):
    """清理和修复常见的JSON格式问题"""
    # 移除可能的前缀说明文字（在第一个[之前的内容）
    start_idx = text.find('[')
    if start_idx > 0:
        text = text[start_idx:]
    
    # 移除可能的后缀说明文字（在最后一个]之后的内容）
    end_idx = text.rfind(']')
    if end_idx > 0 and end_idx < len(text) - 1:
        text = text[:end_idx+1]
    
    # 修复常见的JSON格式错误
    # 1. 单引号替换为双引号（但避免替换已经转义的单引号）
    text = re.sub(r"(?<!\\)'", '"', text)
    
    # 2. 确保属性名有双引号
    text = re.sub(r'([{,])\s*([a-zA-Z0-9_]+)\s*:', r'\1"\2":', text)
    
    # 3. 修复可能的尾随逗号
    text = re.sub(r',\s*([}\]])', r'\1', text)
    
    # 4. 处理JavaScript函数调用，如.repeat()
    text = re.sub(r'\.repeat\(\d+\)', lambda m: '"' + 'a' * 20 + '"', text)  # 替换.repeat(n)为实际重复的字符串
    
    # 5. 处理其他可能的JavaScript函数
    text = re.sub(r'\.[a-zA-Z]+\([^)]*\)', '""', text)  # 将其他JS函数调用替换为空字符串
    
    return text

def generate_case_id(prefix, index=None):
    """生成不重复的测试用例ID"""
    # 使用更精确的时间戳和随机字符生成绝对唯一的ID
    timestamp = datetime.datetime.now().strftime("%Y%m%d%H%M%S")
    # 添加毫秒以增加唯一性
    microseconds = datetime.datetime.now().microsecond // 1000
    
    if index is not None:
        # 使用索引和随机数生成ID
        random_chars = ''.join(random.choices(string.ascii_uppercase + string.digits, k=2))
        return f"{prefix}_{timestamp}_{microseconds:03d}_{index:03d}_{random_chars}"
    else:
        # 使用随机字符生成ID
        random_chars = ''.join(random.choices(string.ascii_uppercase + string.digits, k=4))
        return f"{prefix}_{timestamp}_{microseconds:03d}_{random_chars}"

def get_template_for_type(test_type):
    """根据测试类型返回相应的模板"""
    templates = {
        "功能测试": {
            "fields": ["id", "title", "precondition", "steps", "expected", "priority", "level", "status", "test_data"],
            "example": {
                "id": "F001",
                "title": "功能测试用例标题",
                "precondition": "测试前置条件",
                "steps": ["步骤1", "步骤2"],
                "expected": "预期结果",
                "priority": "P0/P1/P2",  # 改为P格式
                "level": "系统级/集成级/单元级",
                "status": "设计中",
                "test_data": {
                    "username": "test_user",
                    "password": "Test@123",
                    "invalid_inputs": ["", "特殊字符@#$"]
                }
            }
        },
        "性能测试": {
            "fields": ["id", "title", "scenario", "load", "duration", "metrics", "threshold", "level", "expected", "priority", "test_data"],
            "example": {
                "id": "P001",
                "title": "性能测试用例标题",
                "scenario": "测试场景描述",
                "load": "并发用户数/TPS",
                "duration": "测试持续时间",
                "metrics": ["响应时间", "CPU使用率", "内存使用率"],
                "threshold": "性能指标阈值",
                "level": "压力测试/负载测试/容量测试",
                "expected": "预期结果",
                "priority": "P0/P1/P2",  # 添加优先级字段
                "test_data": {
                    "user_count": 1000,
                    "ramp_up_time": "30秒",
                    "test_duration": "10分钟"
                }
            }
        },
        "自动化测试": {
            "fields": ["id", "title", "component", "steps", "expected", "priority", "automation_level", "level", "data_dependency", "test_data"],
            "example": {
                "id": "A001",
                "title": "自动化测试用例标题",
                "component": "测试组件/模块",
                "steps": ["步骤1", "步骤2"],
                "expected": "预期结果",
                "priority": "P0/P1/P2",  # 添加优先级字段
                "automation_level": "完全自动化/半自动化",
                "level": "UI级/API级/单元级",
                "data_dependency": "测试数据依赖",
                "test_data": {
                    "test_env": "测试环境",
                    "browser": "Chrome",
                    "test_account": {
                        "username": "auto_test",
                        "password": "Auto@123"
                    }
                }
            }
        }
    }
    
    return templates.get(test_type, templates["功能测试"])

def generate_with_ollama(data):
    """调用 Ollama 生成测试用例"""
    # 获取测试类型对应的模板
    template = get_template_for_type(data['test_type'])
    fields_str = ", ".join(template["fields"])
    
    # 根据优先级策略设置提示
    priority_hints = {
        "1": "请根据测试用例的实际重要性分配优先级，关键功能和核心场景应为P0（高优先级），次要功能为P1（中优先级），低重要性功能为P2（低优先级）",
        "2": "请将所有测试用例都设为P0（高优先级）",
        "3": "请将所有测试用例都设为P1（中优先级）",
        "4": "请将所有测试用例都设为P2（低优先级）"
    }
    
    priority_hint = priority_hints.get(data.get('priority_strategy', "1"), priority_hints["1"])
    
    # 测试数据生成提示
    test_data_hint = ""
    if data.get('generate_test_data', True):
        test_data_hint = """
    8. 为每个测试用例生成合理的测试数据，包括：
       - 有效和无效的输入值
       - 边界值和极限条件
       - 特殊字符和格式
       - 与测试场景相关的具体数据
    """
    
    prompt = f"""
    请作为资深测试工程师，根据以下场景生成测试用例：
    - 场景名称: {data['name']}
    - 场景描述: {data['description']}
    
    严格要求:
    1. 必须且仅生成 {data['number']} 个{data['test_type']}用例，不多不少，精确为{data['number']}个。
    2. 每个用例必须包含以下字段: {fields_str}
    
    请基于上述场景描述，设计全面覆盖该功能的测试用例，包括正向流程、异常流程、边界条件和特殊场景。测试用例应该能够验证该功能的所有关键方面和可能的用户交互路径。

    要求:
    1. 使用中文描述
    2. 步骤清晰可执行，具体到每个操作步骤
    3. 输出必须是严格的JSON数组格式，不要包含任何其他文本或解释
    4. 请确保JSON格式正确，所有属性名和字符串值都使用双引号
    5. 所有字段必须有有意义的内容，不要留空
    6. {priority_hint}，务必使用P0/P1/P2格式表示优先级（不要使用"高/中/低"）{test_data_hint}
    7. 测试用例应该与场景描述高度相关，确保覆盖描述中提到的所有功能点和业务场景
    8. 用例标题应简洁明了地描述测试目的
    9. 在级别字段中，根据测试类型分别使用：功能测试（系统级/集成级/单元级）、性能测试（压力测试/负载测试/容量测试）、自动化测试（UI级/API级/单元级）
    
    JSON格式示例:
    [
      {{
        {json.dumps(template["example"], ensure_ascii=False, indent=8)[1:-1]}
      }}
    ]
    
    请直接返回JSON数组，不要有任何前缀或后缀说明。记住，必须只生成{data['number']}个测试用例。
    """
    
    try:
        # 使用非流式交互
        response = ollama.chat(
            model='llama3.1:8b',
            messages=[{'role': 'user', 'content': prompt}],
            stream=False  # 改为非流式
        )
        
        # 直接获取完整响应
        content = response.get('message', {}).get('content', '')
        
        # 保存原始响应以便调试
        with open('raw_response.txt', 'w', encoding='utf-8') as f:
            f.write(content)
        
        # 尝试提取JSON部分
        try:
            # 首先尝试直接解析整个响应
            return json.loads(content)
        except json.JSONDecodeError as e:
            # 如果失败，尝试清理和修复JSON
            cleaned_json = clean_json_response(content)
            try:
                return json.loads(cleaned_json)
            except json.JSONDecodeError as e2:
                # 保存响应以便调试
                with open('debug_response.txt', 'w', encoding='utf-8') as f:
                    f.write(content)
                raise Exception("JSON解析失败，详细错误已记录到debug_response.txt")
    
    except Exception as e:
        raise Exception(f"Ollama 调用失败: {str(e)}")

def generate_test_data(test_case, test_type, scenario_name):
    """根据测试用例类型和场景生成合适的测试数据"""
    if test_type == "功能测试":
        if "登录" in scenario_name or "登陆" in scenario_name or "用户" in scenario_name:
            return {
                "valid_username": "test_user",
                "valid_password": "Test@123",
                "invalid_username": ["", "admin'", "<script>alert(1)</script>", "a"*100],
                "invalid_password": ["", "123", "password", "a"*100],
                "special_chars": ["@#$%^&*", "中文用户名", "user@example.com"]
            }
        elif "搜索" in scenario_name:
            return {
                "keywords": ["手机", "电脑", "衣服", "食品"],
                "special_inputs": ["", "*", "SELECT *", "<script>alert(1)</script>"],
                "long_query": "a" * 100,
                "filters": {
                    "price_range": ["100-500", "1000以上"],
                    "category": ["电子产品", "服装"]
                }
            }
        else:
            return {
                "valid_inputs": ["正常输入1", "正常输入2"],
                "invalid_inputs": ["", "特殊字符@#$", "超长输入"+"a"*50],
                "boundary_values": ["0", "1", "999999"]
            }
    
    elif test_type == "性能测试":
        return {
            "user_count": [10, 50, 100, 500, 1000],
            "ramp_up_time": ["10秒", "30秒", "1分钟"],
            "test_duration": ["5分钟", "10分钟", "30分钟"],
            "think_time": ["1-3秒", "3-5秒"],
            "data_volume": ["小(1KB)", "中(1MB)", "大(10MB)"]
        }
    
    elif test_type == "自动化测试":
        return {
            "test_env": ["开发环境", "测试环境", "预发布环境"],
            "browsers": ["Chrome", "Firefox", "Safari", "Edge"],
            "devices": ["PC", "移动设备", "平板"],
            "screen_sizes": ["1920x1080", "1366x768", "375x812"],
            "test_account": {
                "username": "auto_test",
                "password": "Auto@123"
            }
        }
    
    return {}

def post_process_test_cases(test_cases, data):
    """后处理测试用例，添加ID和其他必要信息"""
    # 获取优先级策略
    priority_strategy = data.get('priority_strategy', '1')
    
    # 确保返回的测试用例数量与请求的一致
    requested_count = data.get('number', 1)
    if len(test_cases) > requested_count:
        # 如果生成的用例数量超过了请求的数量，只保留请求的数量
        test_cases = test_cases[:requested_count]
    
    # 处理每个测试用例
    for i, tc in enumerate(test_cases):
        # 如果没有ID，则生成一个
        if not tc.get('id') or tc.get('id') == "":
            tc['id'] = generate_case_id(data.get('id_prefix', 'T'), i+1)
        
        # 确保关键字段不为空
        if not tc.get('title') or tc.get('title') == "":
            tc['title'] = f"{data['name']}测试用例 #{i+1}"
            
        if not tc.get('expected') or tc.get('expected') == "":
            tc['expected'] = f"预期{data['name']}操作成功完成"
        
        # 根据策略设置优先级
        if priority_strategy == "2":  # 全部高优先级
            tc['priority'] = "P0"  # 改为P0格式
        elif priority_strategy == "3":  # 全部中优先级
            tc['priority'] = "P1"  # 改为P1格式
        elif priority_strategy == "4":  # 全部低优先级
            tc['priority'] = "P2"  # 改为P2格式
        # 策略1（基于重要性）保留Ollama生成的优先级，如果没有则设为P1
        elif not tc.get('priority') or tc.get('priority') == "":
            tc['priority'] = "P1"  # 改为P1格式
        # 如果优先级是传统格式，转换为P格式
        elif tc.get('priority') == "高":
            tc['priority'] = "P0"
        elif tc.get('priority') == "中":
            tc['priority'] = "P1"
        elif tc.get('priority') == "低":
            tc['priority'] = "P2"
            
        # 确保所有测试用例类型都有级别字段
        if not tc.get('level') or tc.get('level') == "":
            if data['test_type'] == '功能测试':
                tc['level'] = random.choice(['系统级', '集成级', '单元级'])
            elif data['test_type'] == '性能测试':
                tc['level'] = random.choice(['压力测试', '负载测试', '容量测试'])
            elif data['test_type'] == '自动化测试':
                tc['level'] = random.choice(['UI级', 'API级', '单元级'])
        
        # 添加测试类型
        tc['test_type'] = data['test_type']
            
        # 根据测试类型添加额外字段
        if data['test_type'] == '功能测试':
            if not tc.get('status') or tc.get('status') == "":
                tc['status'] = "设计中"
            if not tc.get('precondition') or tc.get('precondition') == "":
                tc['precondition'] = f"用户已进入{data['name']}页面"
            
        if data['test_type'] == '性能测试':
            if not tc.get('threshold') or tc.get('threshold') == "":
                tc['threshold'] = f"响应时间 < {random.randint(100, 500)}ms"
            if not tc.get('duration') or tc.get('duration') == "":
                tc['duration'] = f"{random.randint(5, 30)}分钟"
            if not tc.get('load') or tc.get('load') == "":
                tc['load'] = f"{random.randint(10, 100)}并发用户"
            
        if data['test_type'] == '自动化测试':
            if not tc.get('automation_level') or tc.get('automation_level') == "":
                tc['automation_level'] = random.choice(['完全自动化', '半自动化'])
            if not tc.get('component') or tc.get('component') == "":
                tc['component'] = f"{data['name']}模块"
            if not tc.get('data_dependency') or tc.get('data_dependency') == "":
                tc['data_dependency'] = "测试数据库"
        
        # 如果需要生成测试数据且没有测试数据
        if data.get('generate_test_data', True) and (not tc.get('test_data') or tc.get('test_data') == ""):
            tc['test_data'] = generate_test_data(tc, data['test_type'], data['name'])
    
    return test_cases

def load_test_cases_from_file(filename='debug_response.txt'):
    """从文件中加载测试用例"""
    try:
        with open(filename, 'r', encoding='utf-8') as f:
            content = f.read()
        return json.loads(content)
    except Exception as e:
        raise Exception(f"从文件加载测试用例失败: {str(e)}")

# 保存测试用例到数据库的函数
def save_test_cases_to_db(test_cases, scenario_name):
    """Save the test cases to the MySQL database"""
    if not DB_AVAILABLE:
        print("数据库连接不可用，测试用例仅保存到文件")
        return {"status": "file_only", "message": "数据库连接不可用，测试用例仅保存到文件"}
    
    successful_saves = 0
    failed_saves = 0
    failed_cases = []
    
    print(f"\n正在将 {len(test_cases)} 个测试用例保存到数据库...")
    
    for tc in test_cases:
        # Map the test case fields to the database schema
        db_test_case = {
            'case_id': tc.get('id', ''),
            'title': tc.get('title', ''),
            'description': tc.get('precondition', '') if tc.get('precondition') else '',
            'content': {
                'steps': tc.get('steps', []),
                'expected_results': tc.get('expected', ''),
                'test_data': tc.get('test_data', {}),
                'level': tc.get('level', ''),
                'precondition': tc.get('precondition', ''),
            }
        }
        
        # Add extra fields based on test type
        if tc.get('threshold'):
            db_test_case['content']['threshold'] = tc.get('threshold')
        if tc.get('duration'):
            db_test_case['content']['duration'] = tc.get('duration')
        if tc.get('load'):
            db_test_case['content']['load'] = tc.get('load')
        if tc.get('automation_level'):
            db_test_case['content']['automation_level'] = tc.get('automation_level')
        if tc.get('component'):
            db_test_case['content']['component'] = tc.get('component')
        if tc.get('data_dependency'):
            db_test_case['content']['data_dependency'] = tc.get('data_dependency')
        
        # Map priority to database format (0=P0 high, 1=P1 medium, 2=P2 low)
        priority_map = {"高": 0, "中": 1, "低": 2}
        db_test_case['priority'] = priority_map.get(tc.get('priority', '中'), 2)
        
        # Map test type to database format
        test_type_map = {
            '功能测试': 1,  # functional
            '性能测试': 4,  # performance
            '自动化测试': 2,  # api (assuming automation tests are API tests)
        }
        db_test_case['test_type'] = test_type_map.get(scenario_name, 1)
        
        # Map status to database format (1=draft, 2=active, 3=archived)
        status_map = {
            '设计中': 1,
            '可执行': 2,
            '已废弃': 3
        }
        db_test_case['status'] = status_map.get(tc.get('status', '设计中'), 1)
        
        # Add tags
        if tc.get('tags'):
            db_test_case['tags'] = tc.get('tags')
        else:
            # Create tags from various fields
            tags = []
            if scenario_name:
                tags.append(scenario_name)
            if tc.get('level'):
                tags.append(tc.get('level'))
            if tc.get('status'):
                tags.append(tc.get('status'))
            db_test_case['tags'] = ','.join(tags)
        
        # Add creator information
        db_test_case['created_by'] = 'Hermes测试用例生成器(Web)'
        
        # Save to database
        success, message, case_id = create_test_case(db_test_case)
        if success:
            successful_saves += 1
        else:
            failed_saves += 1
            failed_cases.append({
                'case_id': tc.get('id', ''),
                'title': tc.get('title', ''),
                'error': message
            })
            print(f"保存用例 {tc.get('id', '')} 失败: {message}")
    
    result = {
        "status": "success" if failed_saves == 0 else "partial",
        "total": len(test_cases),
        "successful": successful_saves,
        "failed": failed_saves,
        "failed_cases": failed_cases if failed_saves > 0 else []
    }
    
    print(f"数据库保存结果: 成功 {successful_saves} 条, 失败 {failed_saves} 条")
    return result

# 路由定义
@app.route('/')
def index():
    """主页"""
    return render_template('index.html')

@app.route('/api/generate', methods=['POST'])
def generate():
    """生成测试用例API"""
    try:
        # 从请求获取参数
        data = {
            "name": request.form.get('name', ''),
            "description": request.form.get('description', ''),
            "number": int(request.form.get('number', 1)),
            "test_type": request.form.get('test_type', '功能测试'),
            "id_prefix": request.form.get('id_prefix', ''),
            "priority_strategy": request.form.get('priority_strategy', '1'),
            "generate_test_data": request.form.get('generate_test_data', 'true').lower() == 'true',
            "save_to_db": request.form.get('save_to_db', 'true').lower() == 'true'
        }
        
        # 验证输入
        if not data['name']:
            return jsonify({"error": "场景名称不能为空"}), 400
        
        if not data['description']:
            data['description'] = data['name']
        
        # 如果未提供前缀，则根据类型设置默认前缀
        if not data['id_prefix']:
            prefix_map = {
                "功能测试": "F", 
                "性能测试": "P",
                "自动化测试": "A"
            }
            data['id_prefix'] = prefix_map.get(data['test_type'], "T")
        
        # 生成测试用例
        test_cases = generate_with_ollama(data)
        
        # 后处理测试用例
        test_cases = post_process_test_cases(test_cases, data)
        
        # 将测试用例保存到本地文件
        timestamp = datetime.datetime.now().strftime("%Y%m%d%H%M%S")
        filename = f"test_cases_{timestamp}.json"
        file_path = os.path.join(app.config['UPLOAD_FOLDER'], filename)
        
        with open(file_path, 'w', encoding='utf-8') as f:
            json.dump(test_cases, f, ensure_ascii=False, indent=2)
        
        result = {
            "success": True,
            "message": "测试用例生成成功",
            "test_cases": test_cases,
            "filename": filename
        }
        
        # 保存到数据库（如果请求要求且数据库可用）
        if data['save_to_db'] and DB_AVAILABLE:
            db_result = save_test_cases_to_db(test_cases, data['test_type'])
            result["db_save"] = db_result
            
        return jsonify(result)
    
    except Exception as e:
        import traceback
        error_traceback = traceback.format_exc()
        
        # 尝试从debug_response.txt加载测试用例
        try:
            test_cases = load_test_cases_from_file()
            
            result = {
                "success": True,
                "message": "使用缓存的测试用例（生成失败）",
                "test_cases": test_cases,
                "error": str(e),
                "traceback": error_traceback
            }
            
            return jsonify(result)
        except:
            return jsonify({
                "success": False,
                "message": "生成测试用例失败",
                "error": str(e),
                "traceback": error_traceback
            }), 500

@app.route('/api/load_from_file', methods=['POST'])
def load_from_file():
    """从文件加载测试用例"""
    try:
        if 'file' not in request.files:
            return jsonify({"error": "没有上传文件"}), 400
            
        file = request.files['file']
        if file.filename == '':
            return jsonify({"error": "没有选择文件"}), 400
            
        # 确保文件名安全
        filename = secure_filename(file.filename)
        filepath = os.path.join(app.config['UPLOAD_FOLDER'], filename)
        file.save(filepath)
        
        # 加载文件内容
        test_cases = load_test_cases_from_file(filepath)
        
        return jsonify({
            "success": True,
            "test_cases": test_cases,
            "count": len(test_cases)
        })
        
    except Exception as e:
        return jsonify({
            "success": False, 
            "error": str(e)
        }), 500

@app.route('/download/<filename>')
def download_file(filename):
    """下载生成的测试用例文件"""
    return send_from_directory(app.config['UPLOAD_FOLDER'], filename, as_attachment=True)

@app.route('/api/search_test_cases', methods=['GET'])
def search_test_cases_api():
    """搜索数据库中的测试用例"""
    if not DB_AVAILABLE:
        return jsonify({
            "success": False,
            "message": "数据库连接不可用"
        }), 503
    
    try:
        # 获取查询参数
        case_id = request.args.get('case_id', '')
        title = request.args.get('title', '')
        test_type = request.args.get('test_type', '')
        status = request.args.get('status', '')
        priority = request.args.get('priority', '')
        tags = request.args.get('tags', '')
        created_by = request.args.get('created_by', '')
        
        # 构建过滤条件
        filters = {}
        if case_id:
            filters['case_id'] = case_id
        if title:
            filters['title'] = title
        if test_type and test_type.isdigit():
            filters['test_type'] = int(test_type)
        if status and status.isdigit():
            filters['status'] = int(status)
        if priority and priority.isdigit():
            filters['priority'] = int(priority)
        if tags:
            filters['tags'] = tags
        if created_by:
            filters['created_by'] = created_by
        
        # 获取分页参数
        limit = min(int(request.args.get('limit', 100)), 100)  # 限制最大返回100条
        offset = int(request.args.get('offset', 0))
        
        # 执行搜索
        success, results = search_test_cases(filters, limit, offset)
        
        if success:
            return jsonify({
                "success": True,
                "count": len(results),
                "test_cases": results
            })
        else:
            return jsonify({
                "success": False,
                "message": results  # 错误消息
            }), 500
    except Exception as e:
        import traceback
        return jsonify({
            "success": False,
            "message": "搜索测试用例失败",
            "error": str(e),
            "traceback": traceback.format_exc()
        }), 500

if __name__ == '__main__':
    app.run(debug=True, host='0.0.0.0', port=5001) 